# Copyright Contributors to the Pyro project.
#
# SPDX-License-Identifier: Apache-2.0

.PHONY: help create-host-workspace build build-gpu run run-gpu notebook notebook-gpu

DOCKER_FILE=Dockerfile
BASE_IMG=ubuntu:24.04
BASE_CUDA_IMG=nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04
DOCKER_CMD=docker
HOST_WORK_DIR=${HOME}/pyro_docker
UID=999
GID=999
OSTYPE=$(shell uname)
USER=pyromancer
DOCKER_WORK_DIR=/home/${USER}/workspace/shared
pyro_git_url=https://github.com/pyro-ppl/pyro.git

# Optional args
python_version?=3.12
pytorch_branch?=release
pyro_branch?=release
cmd?=bash
trust_hosts?="no"

# Determine name of docker image
build run notebook: img_prefix=pyro-cpu
build-gpu run-gpu notebook-gpu: img_prefix=pyro-gpu
build run lab: img_prefix=pyro-cpu
build-gpu run-gpu lab-gpu: img_prefix=pyro-gpu

ifeq ($(img), )
	IMG_NAME=${img_prefix}-${pyro_branch}-${pytorch_branch}-${python_version}
else
	IMG_NAME=${img}
endif

help:
	@fgrep -h "##" ${MAKEFILE_LIST} | fgrep -v fgrep | sed -e 's/##//'

##
##Available targets:
##

build: ##
	## Build a docker image for running Pyro on a CPU.
	## Args:
	##   python_version: version of python to use. default - python 3.12
	##   pytorch_branch: whether to build PyTorch from conda or from source
	##   (git branch specified by pytorch_branch)
	##     default - latest pytorch version on the torch python package index
	##   pyro_branch: whether to use the released Pyro wheel or a git branch.
	##     default - latest pyro-ppl wheel on pypi
	##   trust_hosts: If set to yes hosts SSL ceritificates will be trusted
	##   (might be needed when running begind a firewall)
	##     default - Verify hosts SSL certificates
	##
	${DOCKER_CMD} build -t ${IMG_NAME} \
	--build-arg base_img=${BASE_IMG} \
	--build-arg uid=${UID} \
	--build-arg gid=${GID} \
	--build-arg ostype=${OSTYPE} \
	--build-arg python_version=${python_version} \
	--build-arg pytorch_branch=${pytorch_branch} \
	--build-arg pyro_git_url=${pyro_git_url} \
	--build-arg pyro_branch=${pyro_branch} \
	--build-arg trust_hosts=${trust_hosts} -f ${DOCKER_FILE} .

build-gpu: ##
	## Build a docker image for running Pyro on a GPU.
	## Args:
	##   python_version: version of python to use. default - python 3.12
	##   pytorch_branch: whether to build PyTorch from conda or from source
	##   (git branch specified by pytorch_branch)
	##     default - latest pytorch version on the torch python package index
	##   pyro_branch: whether to use the released Pyro wheel or a git branch.
	##     default - latest pyro-ppl wheel on pypi
	##   trust_hosts: If set to yes hosts SSL ceritificates will be trusted
	##   (might be needed when running begind a firewall)
	##     default - Verify hosts SSL certificates
	##
	${DOCKER_CMD} build -t ${IMG_NAME} \
	--build-arg base_img=${BASE_CUDA_IMG} \
	--build-arg uid=${UID} \
	--build-arg gid=${GID} \
	--build-arg ostype=${OSTYPE} \
	--build-arg pytorch_whl=cu126 \
	--build-arg python_version=${python_version} \
	--build-arg pytorch_branch=${pytorch_branch} \
	--build-arg pyro_git_url=${pyro_git_url} \
	--build-arg pyro_branch=${pyro_branch} \
	--build-arg trust_hosts=${trust_hosts} -f ${DOCKER_FILE} .

create-host-workspace: ##
	## Create shared volume on the host for sharing files with the container.
	##
	mkdir -p ${HOST_WORK_DIR}

run: create-host-workspace
run: ##
	## Start a Pyro CPU docker instance, and run the command `cmd`.
	## Args:
	##   img: use image name given by `img`.
	##   cmd: command invoked on running a docker instance.
	##     default - bash
	##
	docker run --init -it --user ${USER} \
	-v ${HOST_WORK_DIR}:${DOCKER_WORK_DIR} \
	${IMG_NAME} ${cmd}

run-gpu: create-host-workspace
run-gpu: ##
	## Start a Pyro GPU docker instance, and run the command `cmd`.
	## Args:
	##   img: use image name given by `img`.
	##   cmd: command invoked on running a docker instance.
	##     default - bash
	##
	docker run --init --gpus=all -it --user ${USER} \
	-v ${HOST_WORK_DIR}:${DOCKER_WORK_DIR} \
	${IMG_NAME} ${cmd}

notebook: create-host-workspace
notebook: ##
	## Start a jupyter notebook on the Pyro CPU docker container.
	## Args:
	##   img: use image name given by `img`.
	##
	docker run --init -it -p 8888:8888 --user ${USER} \
	-v ${HOST_WORK_DIR}:${DOCKER_WORK_DIR} \
	${IMG_NAME}

notebook-gpu: create-host-workspace
notebook-gpu: ##
	## Start a jupyter notebook on the Pyro GPU docker container.
	## Args:
	##   img: use image name given by `img`.
	##
	docker run --gpus=all --init -it -p 8888:8888 --user ${USER} \
	-v ${HOST_WORK_DIR}:${DOCKER_WORK_DIR} \
	${IMG_NAME}

	notebook: create-host-workspace
lab: ##
	## Start jupyterlab on the Pyro CPU docker container.
	## Args:
	##   img: use image name given by `img`.
	##
	docker run --init -it -p 8888:8888 --user ${USER} \
	-v ${HOST_WORK_DIR}:${DOCKER_WORK_DIR} \
	${IMG_NAME} jupyter lab --port=8888 --no-browser --ip=0.0.0.0

lab-gpu: create-host-workspace
lab-gpu: ##
	## Start jupyterlab on the Pyro GPU docker container.
	## Args:
	##   img: use image name given by `img`.
	##
	docker run --gpus=all --init -it -p 8888:8888 --user ${USER} \
	-v ${HOST_WORK_DIR}:${DOCKER_WORK_DIR} \
	${IMG_NAME} jupyter lab --port=8888 --no-browser --ip=0.0.0.0


